import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn import datasets
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA

if __name__ == '__main__':
    # 加载 iris 数据集
    iris = datasets.load_iris()
    features = iris.data
    labels = iris.target
    # 使用 KMeans 算法进行聚类
    kmeans = KMeans(n_clusters=3, random_state=42)
    kmeans.fit(features)
    cluster_labels = kmeans.labels_
    # 使用 PCA 进行数据降维
    pca = PCA(n_components=2)
    features_pca = pca.fit_transform(features)
    # 可视化原始数据和聚类结果
    plt.figure(figsize=(12, 4))

    plt.subplot(1, 2, 1)
    plt.scatter(features_pca[:, 0], features_pca[:, 1], c=labels, edgecolor='k', s=50)
    plt.title('Original Data (PCA)')
    plt.xlabel('PCA Component 1')
    plt.ylabel('PCA Component 2')

    plt.subplot(1, 2, 2)
    cluster_cmap = ListedColormap(['red', 'green', 'blue'])
    plt.scatter(features_pca[:, 0], features_pca[:, 1], c=cluster_labels, cmap=cluster_cmap, edgecolor='k', s=50)
    plt.title('KMeans Clustering (PCA)')
    plt.xlabel('PCA Component 1')
    plt.ylabel('PCA Component 2')

    plt.tight_layout()
    plt.show()
